import torch.nn as nn
from options import HiDDenConfiguration
from model.conv_bn_relu import ConvBNRelu, ConvBNRelu2
from model.conv_bn_relu import ConvBNSigmoid, ConvBNSigmoidSampling
import torch
import torch.nn.functional as F

class Decoder(nn.Module):
    """
    Decoder module. Receives a watermarked image and extracts the watermark.
    The input image may have various kinds of noise applied to it,
    such as Crop, JpegCompression, and so on. See Noise layers for more.
    """
    def __init__(self, channels, M):
        self.M = M
        super(Decoder, self).__init__()
        self.channels = channels
        self.layers0 = nn.Sequential(ConvBNRelu(3, self.channels))
        self.layers1 = nn.Sequential(ConvBNRelu(self.channels, self.channels))
        self.layers2 = nn.Sequential(ConvBNRelu(2 * self.channels, self.channels))
        self.layers3 = nn.Sequential(ConvBNRelu(3 * self.channels, self.channels))
        self.layers4 = nn.Sequential(ConvBNRelu(self.channels, 30))
        self.layers5 = nn.Sequential(nn.AdaptiveAvgPool2d(output_size=(1, 1)))
        self.linear = nn.Linear(self.M, self.M)
        # self.last_layers2 = nn.Sequential(ConvBNRelu(1, 1))

    # def forward(self, noise_image, image_dct):
    def forward(self, noise_image):

        x1 = self.layers0(noise_image) #(64, 64, 128, 128)


        # print(x1.shape)
        x2 = self.layers1(x1) #(64, 64, 128, 128)
        x2_ = torch.cat([x1, x2], dim=1) #(64, 64*2, 128, 128)

        x3 = self.layers2(x2_) #(64, 64, 128, 128)
        x3_ = torch.cat([x1, x2, x3], dim=1)  #(64, 64*3, 128, 128)

        x4 = self.layers3(x3_) #(64, 64, 128, 128)

        x5 = self.layers4(x4) #(64, 30, 128, 128)

        x = self.layers5(x5) #(64, 30, 1, 1)
        x.squeeze_(3).squeeze_(2) #(64, 30)
        x = self.linear(x) #(64, 30)
        
        return x


# class Decoder(nn.Module):
#     """
#     Decoder module. Receives a watermarked image and extracts the watermark.
#     The input image may have various kinds of noise applied to it,
#     such as Crop, JpegCompression, and so on. See Noise layers for more.
#     """
#     def __init__(self, config: HiDDenConfiguration):
#
#         super(Decoder, self).__init__()
#         self.channels = config.decoder_channels
#         self.layers0 = nn.Sequential(ConvBNRelu(12, self.channels))
#         self.layers1 = nn.Sequential(ConvBNRelu(self.channels, self.channels))
#         self.layers2 = nn.Sequential(ConvBNRelu(2 * self.channels, self.channels))
#         self.layers3 = nn.Sequential(ConvBNRelu(3 * self.channels, self.channels))
#         self.last_layers = nn.Sequential(ConvBNRelu(self.channels, 1))
#
#         # self.layers4 = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
#         # self.last_layers2 = nn.Sequential(ConvBNRelu(1, 1))
#
#     def forward(self, noise_image, image_dct):
#         # x1 = torch.cat([noise_image, image_dct])
#         image_dct = image_dct.reshape([image_dct.size(0), int(3 * 4), int(image_dct.size(2) / 2), int(image_dct.size(2) / 2)])
#         x1 = self.layers0(image_dct)
#
#         x2 = self.layers1(x1)
#         x2_ = torch.cat([x1, x2], dim=1)
#
#         x3 = self.layers2(x2_)
#         x3_ = torch.cat([x1, x2, x3], dim=1)
#
#         x4 = self.layers3(x3_)
#         x = self.last_layers(x4)
#
#         # x = self.layers4(x)
#         # x = self.last_layers2(x)
#         return x